from typing import Self


class Stack:
    arr = []

    def pop(self):
        return self.arr[-1]

    def push(self, e):
        self.arr.append(e)


class Tensor:
    grad: float = 1
    op: str = None
    t1 = None
    t2 = None
    is_leaf = False

    def __init__(self, value: float, grad: float = 1):
        self.value = value
        self.grad = grad

    def backward_v2(self):
        cur = self
        while cur.t1 != None and cur.t1.is_leaf:
            if cur.op == "Add":
                cur.t1.grad = cur.grad
            if cur.op == "Mul":
                cur.t1.grad = cur.grad * cur.t2.value
            cur = cur.t1

    @staticmethod
    def traverse(node):
        if node.t1 != None:
            if node.op == "Add":
                if node.t1.is_leaf:
                    node.t1.grad = node.grad
                if node.t2.is_leaf:
                    node.t2.grad = node.grad
            if node.op == "Mul":
                if node.t1.is_leaf:
                    node.t1.grad = node.grad * node.t2.value
                if node.t2.is_leaf:
                    node.t2.grad = node.grad * node.t1.value
        else:
            return
        Tensor.traverse(node.t1)
        Tensor.traverse(node.t2)


def add(t1: Tensor, t2: Tensor):
    t = Tensor(t1.value + t2.value)
    t.op = "Add"
    t.t1 = t1
    t.t2 = t2
    return t


def mul(t1: Tensor, t2: Tensor):
    t = Tensor(t1.value * t2.value)
    t.op = "Mul"
    t.t1 = t1
    t.t2 = t2
    return t


t1 = Tensor(2)
t2 = Tensor(2.5)
t3 = mul(t1, t2)
t = mul(t3, t2)
t1.is_leaf = True
t2.is_leaf = True
t3.is_leaf = True
# t.grad = 5
# t.backward_v2()
Tensor.traverse(t)
print("t1's grad = %.3f" % t1.grad)
print("t2's grad = %.3f" % t2.grad)
print("t3's grad = %.3f" % t3.grad)
